import os
import re
import pickle
import time
import multiprocessing

from package_.nn_modules import *
from util.draw_process import DrawProcess, save_fig


class CAE(nn.Module):
    """
    CAE: CNN Auto Encoding
    """
    def __init__(self, model_note, in_channels, output_shape, device_):
        """
        :param model_note: 当前模型标签记录，可用于标记当前模型的设计，将模型保存到 model/model_note/ 下。
        :param in_channels: 输入数据的通道数目
        :param output_shape: 输出的大小
        :param device_: cpu or gpu
        """
        self.CONFIG = {
            "model_path": os.path.join("model", str(self.__class__.__name__),  model_note),
            "save_model_weights": False,
            "save_step": 100,
            "save_point": {
                "epoch": 0,
                "val": {
                    "loss": [],
                    "acc": [],
                },
                "train": {
                    "loss": [],
                    "acc": [],
                }
            },
            "save_point_file": os.path.join("model", f"save_{self.__class__.__name__}_point",
                                            model_note, "save_point.pk"),
            'plot_loss_style': {'color': 'red'},
            'plot_acc_style': {'color': 'orange', 'marker': '1'},
            'draw_real_time': False,
        }

        self.in_channels = in_channels
        self.device = device_
        self.y_shape = output_shape
        self.subs_list = [None] * 4

        super(CAE, self).__init__()
        self.nn_stack = nn.Sequential(
            nn.Conv2d(in_channels=self.in_channels, out_channels=32, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Dropout2d(0.25),
            nn.Flatten(),
            nn.Linear(12544, 128),
            nn.ReLU(),
            nn.Dropout2d(0.5),
            nn.Linear(128, self.y_shape),
            nn.LogSoftmax(dim=1)
        )

    def forward(self, x):
        output = self.nn_stack(x)
        return output

    def cae_train(self, dataloader, loss_fn, optimizer):
        size = len(dataloader.dataset)
        time_start = time.time()
        print(f"Epoch {self.CONFIG['save_point']['epoch'] + 1}\n===============================================")
        acc_final = 0
        for batch, (X, y) in enumerate(dataloader):
            X, y = X.to(self.device), y.to(self.device)

            pred = self(X)
            loss = loss_fn(pred, y)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            if (batch + 1) % self.CONFIG['save_step'] == 0:
                loss, current = loss.item() / len(X), batch * len(X)
                acc = 100 * (pred.argmax(1) == y).type(torch.float).sum().item() / len(X)
                acc_final += acc
                self._add_train_point(loss, acc)
                print(f"loss: {loss:>7f}    acc: {acc:>0.1f}%   [{current:>5d}/{size:>5d}]  eplapsed: {(time.time() - time_start):>0.6f}S")
                time_start = time.time()
                if self.CONFIG['draw_real_time']:
                    self._draw_real_time()
        # if self.CONFIG['save_model_weights']:
        #     self._save_cae_weights(f'{(acc_final / batch):>0.1f}')

    def cae_val(self, dataloader, loss_fn, val_or_test=True):
        size = len(dataloader.dataset)
        test_loss, correct = 0, 0
        self.eval()
        with torch.no_grad():
            for X, y in dataloader:
                X, y = X.to(self.device), y.to(self.device)
                pred = self(X)
                test_loss += loss_fn(pred, y).item()
                correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        test_loss /= size
        correct /= size
        self._add_test_point(test_loss, 100 * correct)
        if val_or_test:
            print(f"Validation: \nAccuracy: {(100 * correct):>0.1f}%, Avg Loss: {test_loss:>7f} \n")
            if self.CONFIG['save_model_weights']:
                self._save_cae_weights(f'{correct:>0.1f}')
        else:
            print(f"Test: \nAccuracy: {(100 * correct):>0.1f}%, Avg Loss: {test_loss:>7f} \n")

    def cae_test(self, dataloader, loss_fn):
        self.cae_val(dataloader, loss_fn, False)
        # size = len(dataloader.dataset)
        # self.eval()
        # test_loss, correct = 0, 0

        # with torch.no_grad():
        #     for X, y in dataloader:
        #         X, y = X.to(self.device), y.to(self.device)
        #         pred = self(X)
        #         test_loss += loss_fn(pred, y).item()
        #         correct += (pred.argmax(1) == y).type(torch.float).sum().item()

        # test_loss /= size
        # correct /= size
        # self._add_test_point(test_loss, 100 * correct)
        # print(f"Test: \nAccuracy: {(100 * correct):>0.1f}%, Avg Loss: {test_loss:>7f} \n")

    def predict(self, X):
        """
        :param X: 输入数据 X
        :return: 预测值 y
        """
        # model.eval() pytorch 会自动把 BN 和 DropOut 固定住，不会取平均，而是用训练好的值。
        self.eval()
        return self(X)

    def pre_train(self, epoch):
        """
        加载之前的模型参数
        :return: None
        """
        last_version = self._get_save_model_file(epoch)
        if last_version is None:
            raise print("Error! no model found.")
        self.load_state_dict(torch.load(last_version))
        # self.eval()

        with open(self.CONFIG['save_point_file'], "rb") as f:
            self.CONFIG['save_point'] = pickle.load(f)

        print("pretrain over.")

    def is_save_weights(self, flag):
        """
        设置 flag=True，训练过程中自动保存模型参数。
        :param flag: True or False
        :return: None
        """
        self.CONFIG['save_model_weights'] = flag

    def load_point(self):
        """
        :return: 包含：epoch, loss，acc 的字典，其中 epoch 为 int, loss 和 acc 为 list
        """
        with open(self.CONFIG['save_point_file'], "rb") as f:
            save_point = pickle.load(f)
        return save_point

    def _add_test_point(self, loss, acc):
        self.CONFIG['save_point']['val']['loss'].append(loss)
        self.CONFIG['save_point']['val']['acc'].append(acc)
        self.CONFIG['save_point']['epoch'] += 1
        if not os.path.exists(self.CONFIG['save_point_file']):
            os.makedirs(os.path.dirname(self.CONFIG['save_point_file']))

        with open(self.CONFIG['save_point_file'], "wb") as f:
            pickle.dump(self.CONFIG['save_point'], f)

    def _add_train_point(self, loss, acc):
        self.CONFIG['save_point']['train']['loss'].append(loss)
        self.CONFIG['save_point']['train']['acc'].append(acc)

    def _get_save_model_file(self, epoch):
        last_version = os.path.join(self.CONFIG['model_path'], f'{self.__class__.__name__}-epoch-{epoch}.pth')
        # if len(last_version) == 0:
        #     return None
        # elif self.__class__.__name__ in last_version[0]:
        #     return os.path.join(self.CONFIG['model_path'], last_version[0])
        return last_version

    def _save_cae_weights(self, accuracy):
        if not os.path.exists(self.CONFIG['model_path']):
            os.mkdir(self.CONFIG['model_path'])
        # last_version = self._get_save_model_file()
        # if last_version is not None:
            # 获取已保存模型的正确率
            # last_acc = float(".".join(re.findall(r"\d+", last_version)))
            # last_acc = float(last_version.split(os.sep)[-1].split("-")[-1].replace("_", ".").replace(".pth", ""))
            # if last_acc > float(accuracy):  # 保存正确率高的模型 (训练集,一定程度上减少存储次数)
            #     return
            # else:
            #     os.remove(last_version)
            # os.remove(last_version)
        
        torch.save(self.state_dict(),
                    os.path.join(self.CONFIG['model_path'], f"{self.__class__.__name__}-epoch-{self.CONFIG['save_point']['epoch']}.pth"))
                #    f"{self.CONFIG['model_path']}/{self.__class__.__name__}-{'_'.join(accuracy.split('.'))}.{self.CONFIG['save_point']['epoch']}.pth")

    def _draw_real_time(self):
        # 在multiprocessing的manager共享内存中，除了一级赋值操作，其他的都无效。如[0]["data"]的赋值操作为二级，无效；append非赋值，无效
        # self.subs_list[0]["data"].append(self.CONFIG["save_point"]["train"]["loss"][-1])
        # self.subs_list[1]["data"].append(self.CONFIG["save_point"]["train"]["acc"][-1])
        # self.subs_list[2]["data"].append(self.CONFIG["save_point"]["test"]["loss"][-1])
        # self.subs_list[3]["data"].append(self.CONFIG["save_point"]["test"]["acc"][-1])
        self._write_to_sub_process()

    def _write_to_sub_process(self):
        self.lock.acquire()
        self.build_subs_list()
        self.lock.release()

    def build_subs_list(self):
        self.subs_list[0] = {"title": "train: epoch - loss", "xlabel": "epoch", "ylabel": "loss",
                             "data": self.CONFIG['save_point']['train']['loss'],
                             "style": self.CONFIG['plot_loss_style']}
        self.subs_list[1] = {"title": "train: epoch - accuracy", "xlabel": "epoch", "ylabel": "accuracy",
                             "data": self.CONFIG['save_point']['train']['acc'],
                             "style": self.CONFIG['plot_acc_style']}
        self.subs_list[2] = {"title": "validation: epoch - loss", "xlabel": "epoch", "ylabel": "loss",
                             "data": self.CONFIG['save_point']['val']['loss'],
                             "style": self.CONFIG['plot_loss_style']}
        self.subs_list[3] = {"title": "validation: epoch - accuracy", "xlabel": "epoch", "ylabel": "accuracy",
                             "data": self.CONFIG['save_point']['val']['acc'],
                             "style": self.CONFIG['plot_acc_style']}

    def set_draw_real_time(self):
        self.CONFIG['draw_real_time'] = True

        manager = multiprocessing.Manager()
        self.lock = manager.Lock()
        self._write_to_sub_process()
        self.subs_list = manager.list(self.subs_list)

        proc = multiprocessing.Process
        self.dp = proc(target=DrawProcess, args=("22", self.lock, self.subs_list, "EPOCH -- LOSS & ACCURACY"))
        self.dp.start()

    def save_process(self, title, path):
        self.CONFIG['save_point'] = self.load_point()
        self.build_subs_list()
        save_fig(self.subs_list, title, path)

    def draw_process(self):
        self.CONFIG['save_point'] = self.load_point()
        self.set_draw_real_time()
        self._draw_real_time()
        self.dp.join()


if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Using {device} device")
    cae = CAE(3, 10, device).to(device)
